#include <Lucnx/memory.h>
#include <Lucnx/printk.h>
#include <Lucnx/mutex.h>
#include <Lucnx/system.h>
#include <Lucnx/string.h>
#include <Lucnx/stdlib.h>
#include <Lucnx/debug.h>
#include <Lucnx/task.h>

#define ZONE_VALID 1    // ards 可用内存区域

#define KERNEL_VADDR_START 0x0
#define BITMAP_HEAP_ADDR 0x9FBFF

#define PDE_IDX(vaddr) ((vaddr & 0xffc00000) >> 22)
#define PTE_IDX(vaddr) ((vaddr & 0x003ff000) >> 12)

#define IS_PAGE(address) !(address & 0xfff)

typedef struct memory_pool_t
{
    bitmap_t bitmap;
    lock_t lock;
    u32 memory_start;
}memory_pool_t;

extern void karena_init();

extern sysinfo_t *system_info;
static memory_pool_t kernel_pool;
static memory_pool_t user_pool;

u32 memory_size;

static void mem_init(u32 magic, u32 addr)
{
    u32 memory_base;
    u32 ards_count = *((u32*)addr);
    ards_t *ptr = (ards_t *) (addr + 4);

    for (size_t i = 0; i < ards_count; i++, ptr++)
    {
        if (ptr->length > memory_size && ptr->type == ZONE_VALID)
        {
            memory_size = (u32)ptr->length + (u32)ptr->base;
        }
    }
    assert(memory_size / 1024 / 1024 > 16); // 内存至少 16 MB

    u32 kernel_mem_size = KERNEL_MEM_SIZE; // 内核只占用10MB, 不会因为内存大小而改变
    u32 kernel_pool_size = KERNEL_MEM_SIZE - 0x120000;
    u32 user_mem_size = (memory_size - kernel_mem_size) / PAGE_SIZE * PAGE_SIZE;

    kernel_pool.memory_start = 0x120000;
    user_pool.memory_start = kernel_mem_size;

    u32 kbm_len = kernel_pool_size / PAGE_SIZE / 8;
    kernel_pool.bitmap.bits = (u8 *) BITMAP_HEAP_ADDR - kbm_len;
    kernel_pool.bitmap.btmp_len = kbm_len;

    u32 ubm_len = user_mem_size / PAGE_SIZE / 8;
    user_pool.bitmap.bits = (u8 *) BITMAP_HEAP_ADDR - kbm_len - ubm_len;
    user_pool.bitmap.btmp_len = ubm_len;

    bitmap_init(&kernel_pool.bitmap);
    bitmap_init(&user_pool.bitmap);

    lock_init(&kernel_pool.lock);
    lock_init(&user_pool.lock);

    INFOK("INFO", "Memory init done, memory size: %dB", memory_size);
}

#define PAGE_DIR_ADDR 0x1000

void page_init()
{
    u32 *pde = (u32*)PAGE_DIR_ADDR;
    memset(pde, 0, PAGE_SIZE * 2);

    // 映射前 2 MB 内存
    pde[0] = (0x2000) | 7;
    pde[1] = (0x3000) | 7;

    // 把最后一项设置为页目录自己, 方便修改
    pde[1023] = PAGE_DIR_ADDR | 7;

    u32 *pte = (u32 *) 0x2000;
    for (size_t i = 0; i < 1024; i++)
    {
        pte[i] = (i * PAGE_SIZE) | 7;
    }

    pte = (u32 *) 0x3000;
    for (size_t i = 0; i < 1024; i++)
    {
        pte[i] = (i * PAGE_SIZE + 0x400000) | 7;
    }

    // 开启分页机制
    asm volatile(
        "movl $0x1000, %eax\n"
        "movl %eax, %cr3\n"
        "movl %cr0, %eax\n"
        "orl $0x80000000, %eax\n"
        "movl %eax, %cr0\n");
}

void memory_init(u32 magic, u32 addr)
{
    mem_init(magic, addr);
    page_init();
    karena_init();
}

void memory_lock_acquire(pool_flag pf)
{
    lock_t *lock = pf ? &user_pool.lock : &kernel_pool.lock;
    lock_acquire(lock);
}

void memory_lock_release(pool_flag pf)
{
    lock_t *lock = pf ? &user_pool.lock : &kernel_pool.lock;
    lock_release(lock);
}

// 获取虚拟地址对应的 PDE 指针
static u32 *pde_ptr(void *vaddr)
{
    u32 ptr = 0xFFFFF000 + (PDE_IDX((u32)vaddr) * 4);
    return (u32*)ptr;
}

// 获取虚拟地址对应的 PTE 指针
static u32 *pte_ptr(void* vaddr)
{
    u32 _vaddr = (u32)vaddr;
    u32 ptr = 0xFFC00000 + (PDE_IDX(_vaddr) << 12) + (PTE_IDX(_vaddr) * 4);
    return (u32*)ptr;
}

// 虚拟地址转化为物理地址
void *addr_v2p(void *vaddr)
{
    u32 *pte = pte_ptr(vaddr);
    return (void*) ((*pte & 0xfffff000) + ((u32)vaddr & 0xfff));
}

// 分配 count 页物理内存
static void *page_alloc(pool_flag pf, int count)
{
    memory_pool_t *pool = pf ? &user_pool : &kernel_pool;
    int idx = bitmap_scan(&pool->bitmap, count);
    if (idx == EOF)
    {
        return NULL;
    }
    for (size_t i = 0; i < count; i++)
    {
        bitmap_set(&pool->bitmap, idx + i, 1);
    }
    return (void*) (pool->memory_start + idx * PAGE_SIZE);
}

// 释放 count 页物理内存
void page_free(void *page, int count, pool_flag pf)
{
    memory_pool_t *pool = pf ? &user_pool : &kernel_pool;
    u32 _page = (u32)page; 

    int idx = (_page - pool->memory_start) / PAGE_SIZE;
    for (size_t i = 0; i < count; i++)
    {
        bitmap_set(&pool->bitmap, idx + i, 0);
    }
}

// 映射
void page_map(void *vaddr, void *paddr)
{
    u32 _vaddr = (u32)vaddr;
    u32 _paddr = (u32)paddr;
    assert(IS_PAGE(_paddr) && IS_PAGE(_vaddr));

    u32 *pte = pte_ptr(vaddr);
    u32 *pde = pde_ptr(vaddr);

    if (*pde & 1)
    {
        if (*pte & 1)
        { 
            return;
        }
        *pte = _paddr | 7;
    }
    else
    {
        *pde = (u32) page_alloc(PF_KERNEL, 1) | 7;
        memset((void*)((u32)pte & 0xfffff000), 0, PAGE_SIZE);

        if (*pte & 1)
        {
            panic("Memory mapping vaddr: %p pte: %p pde:%p\n", vaddr, pte, pde);
        }
        *pte = _paddr | 7;
    }
}

void put_page(void *vpage)
{
    u32 *pte = pte_ptr(vpage);
    *pte &= 0xFFFFFFFE;
    asm volatile("invlpg %0" :: "m"(vpage) : "memory");
}

void *page_link2(void *vpage, task_t *task)
{
    u32 _vpage = (u32)vpage;
    if (_vpage < USER_MEMORY)
    {
        printk("kernel memory mapped!\n");
        return NULL;
    }
    void *page = page_alloc(PF_USER, 1);
    if (page == NULL)
    {
        printk("The memory has been allocated\n");
    }
    bitmap_t *map = &task->vmap;

    int pidx = (_vpage - USER_MEMORY) >> 12;
    assert(!bitmap_scan_test(map, pidx));
    bitmap_set(map, pidx, 1);
    page_map(vpage, page);
}

void *page_link(void *vpage)
{
    page_link2(vpage, running_task());
}

// 取消映射
void page_unlink(void *vpage)
{
    u32 page = (u32)addr_v2p(vpage);

    u32 _vpage = (u32)vpage;
    if (_vpage > USER_MEMORY && _vpage < USER_STACK)
    {
        task_t *task = running_task();
        int pidx = (_vpage - USER_MEMORY) >> 12;
        assert(bitmap_scan_test(&task->vmap, pidx));
        bitmap_set(&task->vmap, pidx, 0);
    }

    memory_pool_t *pool= page > user_pool.memory_start ? &user_pool : &kernel_pool;
    int pidx = (page - pool->memory_start) >> 12;
    bitmap_set(&pool->bitmap, pidx, 0);
    put_page(vpage);

    LOGK("Unlink a page %p", vpage);
}

void *kpage_alloc(int count)
{
    memory_lock_acquire(PF_KERNEL);
    void *page = page_alloc(PF_KERNEL, count);
    memory_lock_release(PF_KERNEL);

    INFOK("LOGK", "Alloc %d kernel page%c %p", count, count == 1 ? ' ' : 's', page);
    return page;
}

void kpage_free(void *page, int count)
{
    memory_lock_acquire(PF_KERNEL);
    page_free(page, count, PF_KERNEL);
    memory_lock_release(PF_KERNEL);

    INFOK("LOGK", "Free %d kernel page%c %p", count, count > 1 ? 's' : ':', page);
}

int sys_brk(void *brk)
{
    task_t *task = running_task();
    u32 page = (u32) brk;
    assert(task->mode == 2);
    assert(IS_PAGE(page));

    if (page >= USER_STACK || page < USER_MEMORY)
    {
        printk("brk: out of memory!\n");
        return EOF;
    }

    u32 old_brk = task->brk;

    if (page < task->brk)
    {
        for (size_t i = 0; i < (old_brk - page) >> 12; i++)
        {
            page_unlink((void *) (page + i * PAGE_SIZE));
        }
    }
    else
    {
        for (size_t i = 0; i < (page - old_brk) >> 12; i++)
        {
            page_link((void*) (old_brk + i * PAGE_SIZE));
        }
    }
    
    task->brk = brk;
    return 0;
}

void build_task_stack3()
{
    void *stack = (void *)USER_STACK;

    for (size_t i = 0; i < 4; i++)
    {
        void *page = page_alloc(PF_USER, 1);
        page_map(stack, page);
        stack += PAGE_SIZE;
    }
}

void release_task_stack3()
{
    void *stack = (void *)USER_STACK;

    for (size_t i = 0; i < 4; i++)
    {
        void *page = addr_v2p(stack);
        page_free(page, 1, PF_KERNEL);
        put_page(stack);
        stack += PAGE_SIZE;
    }
}

void copy_task_mem(task_t *task)
{
    task_t *parent = running_task();
    bitmap_t *map = &parent->vmap;
    void *kbuf = kpage_alloc(1);
    for (int pidx = 0; pidx < map->btmp_len * 8; pidx++)
    {
        if (!bitmap_scan_test(map, pidx))
        {
            continue;
        }
        void *vpage = (void *) (USER_MEMORY + pidx * PAGE_SIZE);
        memcpy(kbuf, vpage, PAGE_SIZE);
        task_activate(task);
        page_link2(vpage, task);
        memcpy(vpage, kbuf, PAGE_SIZE);
        task_activate(parent);
    }
    kpage_free(kbuf, 1);
    kbuf = kpage_alloc(4);
    memcpy(kbuf, (void *) USER_STACK, PAGE_SIZE * 4);
    task_activate(task);
    build_task_stack3();
    memcpy((void *)USER_STACK, kbuf, PAGE_SIZE * 4);
    task_activate(parent);
    kpage_free(kbuf, 4);
}

void free_task_mem()
{
    task_t *task = running_task();
    bitmap_t *map = &task->vmap;

    for (size_t i = 0; i < map->btmp_len * 8; i++)
    {
        if (!bitmap_scan_test(map, i))
        {
            continue;
        }
        // 取消映射 释放物理地址
        void *vpage = (void *) (USER_MEMORY + i * PAGE_SIZE);
        page_unlink(vpage);
    }
    // 释放三级栈
    release_task_stack3();
}

void free_pde()
{
    kpage_free(running_task()->page_dir, 1);
}

void free_vmap()
{
    task_t *task = running_task();
    bitmap_t *map = &task->vmap;
    kpage_free(map->bits, 1);
}